package com.ah;

import javassist.*;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

public class TestTransformer implements ClassFileTransformer {
    private final String CLASS_NAME = "com.ah.controller.HelloController";

    private final String METHOD_NAME = "hello";

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        //将所有的字节码的\替换成.
        className = className.replaceAll("/",".");
//        System.out.println("className = " + className);
        if(CLASS_NAME.equals(className)){
            System.out.println("className = " + className);
            //获取字节码
            try {
                CtClass ctClass = ClassPool.getDefault().get(CLASS_NAME);
                //获取字节码的方法
                CtMethod helloMethod = ctClass.getDeclaredMethod(METHOD_NAME);

                //方法增强,方法执行之前
                helloMethod.insertBefore("long begin = System.currentTimeMillis();");
                //方法增强,方法执行之后
                helloMethod.insertAfter("long end = System.currentTimeMillis();");
                helloMethod.insertAfter("System.out.println(\"耗时 = \" + (end-begin));");

                //将增强后的字节码返回
                return ctClass.toBytecode();
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }

        return null;
    }

}
